import sys
import cv2
import face_recognition
import numpy as np
import os
import csv
import datetime
from PyQt5.QtWidgets import (
    QWidget,
    QLabel,
    QVBoxLayout,
    QHBoxLayout,
    QPushButton,
    QLineEdit,
    QTableWidget,
    QTableWidgetItem,
    QApplication,
)
from PyQt5.QtCore import QTimer, Qt
from PyQt5.QtGui import QImage, QPixmap
from imgaug import augmenters as iaa  # 数据增强

class FaceCaptureApp(QWidget):
    def __init__(self):
        super().__init__()
        self.cap = cv2.VideoCapture(0)  # 默认摄像头
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

        # 初始化 Dlib 的人脸编码
        self.registered_encodings = {}
        self.load_registered_students()

        # 界面布局
        self.init_ui()

        # 定时器：每30ms更新画面
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.update_frame)
        self.timer.start(30)

    def init_ui(self):
        self.setWindowTitle("智能人脸识别考勤系统")
        self.setGeometry(100, 100, 800, 600)

        # 主布局
        main_layout = QVBoxLayout()

        # 视频显示区域
        self.video_label = QLabel(self)
        self.video_label.setFixedSize(640, 480)
        main_layout.addWidget(self.video_label)

        # 控制面板（底部）
        control_layout = QHBoxLayout()

        # 注册学生输入
        self.register_id_input = QLineEdit()
        self.register_id_input.setPlaceholderText("输入学号")
        self.register_name_input = QLineEdit()
        self.register_name_input.setPlaceholderText("输入姓名")

        # 按钮
        self.register_btn = QPushButton("注册学生")
        self.register_btn.clicked.connect(self.register_student)
        self.clear_btn = QPushButton("清除考勤")
        self.clear_btn.clicked.connect(self.clear_attendance)

        control_layout.addWidget(self.register_id_input)
        control_layout.addWidget(self.register_name_input)
        control_layout.addWidget(self.register_btn)
        control_layout.addWidget(self.clear_btn)

        # 考勤记录表格
        self.record_table = QTableWidget()
        self.record_table.setColumnCount(3)
        self.record_table.setHorizontalHeaderLabels(["时间", "学号", "姓名"])
        main_layout.addWidget(self.record_table)

        # 状态标签
        self.status_label = QLabel("准备就绪")
        self.status_label.setAlignment(Qt.AlignBottom | Qt.AlignLeft)
        main_layout.addWidget(self.status_label)

        # 整合布局
        self.setLayout(main_layout)

    def load_registered_students(self):
        """加载已注册的学生人脸编码"""
        self.registered_encodings = {}
        registered_dir = "registered_faces"
        if not os.path.exists(registered_dir):
            os.makedirs(registered_dir)
            return

        for student_id in os.listdir(registered_dir):
            student_path = os.path.join(registered_dir, student_id)
            if not os.path.isdir(student_path):
                continue

            # 加载所有图像并计算编码
            encodings = []
            for filename in os.listdir(student_path):
                if not filename.endswith((".jpg", ".png")):
                    continue
                image_path = os.path.join(student_path, filename)
                image = face_recognition.load_image_file(image_path)
                try:
                    encoding = face_recognition.face_encodings(image)[0]
                    encodings.append(encoding)
                except IndexError:
                    continue  # 无检测到人脸的图像忽略

            if encodings:
                self.registered_encodings[student_id] = np.mean(encodings, axis=0)

    def register_student(self):
        """注册新学生（含数据增强）"""
        student_id = self.register_id_input.text().strip()
        student_name = self.register_name_input.text().strip()
        if not student_id or not student_name:
            self.show_status("请输入学号和姓名")
            return

        # 创建学生文件夹
        student_dir = os.path.join("registered_faces", student_id)
        if not os.path.exists(student_dir):
            os.makedirs(student_dir)
        else:
            self.show_status("该学号已存在")
            return

        # 采集当前帧作为注册图像
        ret, frame = self.cap.read()
        if not ret:
            self.show_status("无法获取摄像头图像")
            return

        # 保存原始图像
        original_path = os.path.join(student_dir, f"{student_id}_original.jpg")
        cv2.imwrite(original_path, frame)

        # 数据增强并保存
        self.augment_and_save_images(original_path, student_dir, 10)

        # 更新注册编码
        self.load_registered_students()
        self.show_status(f"注册成功：{student_id} {student_name}", 5000)

    def augment_and_save_images(self, original_path, output_dir, count=10):
        """图像增强：生成 count 张变换后的图像"""
        image = cv2.imread(original_path)
        seq = iaa.Sequential([
            iaa.Fliplr(0.5),          # 水平翻转
            iaa.Affine(rotate=(-20, 20)),  # 旋转 ±20 度
            iaa.GaussianBlur(sigma=(0, 3.0)),  # 高斯模糊
            iaa.Add((-50, 50)),       # 亮度调整
        ])

        for i in range(count):
            augmented = seq.augment_image(image)
            cv2.imwrite(f"{output_dir}/{student_id}_{i}.jpg", augmented)

    def update_frame(self):
        """实时检测和显示"""
        ret, frame = self.cap.read()
        if not ret:
            return

        # 转为 RGB 格式（Dlib 需要）
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        # 预处理：直方图均衡化 + 高斯模糊
        processed_frame = self.process_frame(rgb_frame)

        # 检测人脸
        face_locations = face_recognition.face_locations(processed_frame)
        face_encodings = face_recognition.face_encodings(processed_frame, face_locations)

        # 绘制检测框
        for (top, right, bottom, left), face_encoding in zip(face_locations, face_encodings):
            matches = face_recognition.compare_faces(
                list(self.registered_encodings.values()),
                face_encoding,
                tolerance=0.5  # 调整阈值（0.4-0.6）
            )

            if True in matches:
                student_id = list(self.registered_encodings.keys())[matches.index(True)]
                self.show_status(f"识别成功：{student_id}", 2000)
                self.record_attendance(student_id)
            else:
                self.show_status("未识别到注册人脸", 2000)

            # 绘制矩形框
            cv2.rectangle(frame, (left, top), (right, bottom), (0, 255, 0), 2)

        # 显示到 video_label
        h, w, ch = frame.shape
        bytes_per_line = ch * w
        qImg = QImage(
            frame.data, w, h, bytes_per_line, QImage.Format_RGB888BGR
        ).rgbSwapped()
        self.video_label.setPixmap(QPixmap.fromImage(qImg))

    def process_frame(self, frame):
        """图像预处理：直方图均衡化 + 高斯模糊"""
        gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)
        gray = cv2.GaussianBlur(gray, (5,5), 0)
        return gray

    def record_attendance(self, student_id):
        """记录考勤到表格和CSV"""
        current_time = datetime.datetime.now().strftime("%H:%M:%S")
        student_name = "未知"  # 假设从数据库获取姓名（此处简化）

        # 更新表格
        row_count = self.record_table.rowCount()
        self.record_table.insertRow(row_count)
        self.record_table.setItem(row_count, 0, QTableWidgetItem(current_time))
        self.record_table.setItem(row_count, 1, QTableWidgetItem(student_id))
        self.record_table.setItem(row_count, 2, QTableWidgetItem(student_name))

        # 保存到CSV
        with open("attendance.csv", "a", newline="") as f:
            writer = csv.writer(f)
            writer.writerow([current_time, student_id, student_name])

    def show_status(self, message, duration=3000):
        """更新状态栏"""
        self.status_label.setText(f"状态：{message}")
        QTimer.singleShot(duration, lambda: self.status_label.setText("准备就绪"))

    def clear_attendance(self):
        """清空考勤记录"""
        self.record_table.setRowCount(0)
        self.show_status("考勤记录已清空", 2000)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = FaceCaptureApp()
    window.show()
    sys.exit(app.exec_())